{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMOLgidIE0joOlL2VuTcLPa",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/relax_amber.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#relax your structure (using amber)"
      ],
      "metadata": {
        "id": "TXSecRRnpGeN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title install amber\n",
        "from IPython.utils import io\n",
        "import os\n",
        "import subprocess\n",
        "import tqdm.notebook\n",
        "from sys import version_info \n",
        "python_version = f\"{version_info.major}.{version_info.minor}\"\n",
        "\n",
        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
        "\n",
        "if not os.path.isfile(\"stereo_chemical_props.txt\"):\n",
        "  try:\n",
        "    with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      with io.capture_output() as captured:\n",
        "        # Install py3dmol.\n",
        "        %shell pip install py3dmol\n",
        "        pbar.update(4)\n",
        "\n",
        "        # Install OpenMM and pdbfixer.\n",
        "        %shell rm -rf /opt/conda\n",
        "        %shell wget -q -P /tmp \\\n",
        "          https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \\\n",
        "            && bash /tmp/Miniconda3-latest-Linux-x86_64.sh -b -p /opt/conda \\\n",
        "            && rm /tmp/Miniconda3-latest-Linux-x86_64.sh\n",
        "        pbar.update(15)\n",
        "\n",
        "        PATH=%env PATH\n",
        "        %env PATH=/opt/conda/bin:{PATH}\n",
        "        %shell conda install -qy conda==4.13.0 \\\n",
        "            && conda install -qy -c conda-forge \\\n",
        "              python={python_version} \\\n",
        "              openmm=7.5.1 \\\n",
        "              pdbfixer\n",
        "        pbar.update(80)\n",
        "\n",
        "        %shell wget -q -P /content \\\n",
        "          https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt\n",
        "        pbar.update(1)\n",
        "  except subprocess.CalledProcessError:\n",
        "    print(captured)\n",
        "    raise"
      ],
      "metadata": {
        "cellView": "form",
        "id": "1KKjet38pMys"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title install AlphaFold\n",
        "import sys\n",
        "if not os.path.isdir(\"/content/alphafold\"):\n",
        "  GIT_REPO = 'https://github.com/deepmind/alphafold'\n",
        "  try:\n",
        "    with tqdm.notebook.tqdm(total=100, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
        "      with io.capture_output() as captured:\n",
        "        %shell rm -rf alphafold\n",
        "        %shell git clone --branch main {GIT_REPO} alphafold\n",
        "        pbar.update(10)\n",
        "        # Install the required versions of all dependencies.\n",
        "        %shell pip3 install -r ./alphafold/requirements.txt\n",
        "        # Run setup.py to install only AlphaFold.\n",
        "        %shell pip3 install --no-dependencies ./alphafold\n",
        "        pbar.update(90)\n",
        "\n",
        "        # Apply OpenMM patch.\n",
        "        %shell pushd /opt/conda/lib/python{python_version}/site-packages/ && \\\n",
        "            patch -p0 < /content/alphafold/docker/openmm.patch && \\\n",
        "            popd\n",
        "\n",
        "        # Make sure stereo_chemical_props.txt is in all locations where it could be searched for.\n",
        "        %shell mkdir -p /content/alphafold/alphafold/common\n",
        "        %shell cp -f /content/stereo_chemical_props.txt /content/alphafold/alphafold/common\n",
        "        %shell mkdir -p /opt/conda/lib/python{python_version}/site-packages/alphafold/common/\n",
        "        %shell cp -f /content/stereo_chemical_props.txt /opt/conda/lib/python{python_version}/site-packages/alphafold/common/\n",
        "\n",
        "  except subprocess.CalledProcessError:\n",
        "    print(captured)\n",
        "    raise\n",
        "\n",
        "if \"/content/alphafold\" not in sys.path:\n",
        "  # Make sure everything we need is on the path.\n",
        "  sys.path.append(f\"/opt/conda/lib/python{python_version}/site-packages\")\n",
        "  sys.path.append('/content/alphafold')\n",
        "\n",
        "import jax\n",
        "if jax.local_devices()[0].platform == 'cpu':\n",
        "  DEVICE = \"CPU\"\n",
        "else:\n",
        "  DEVICE = \"GPU\"\n",
        "print(f'Will run on {DEVICE}')\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
        "\n",
        "from alphafold.relax import relax\n",
        "from alphafold.relax import utils\n",
        "from alphafold.common import protein, residue_constants\n",
        "\n",
        "MODRES = {'MSE':'MET','MLY':'LYS','FME':'MET','HYP':'PRO',\n",
        "          'TPO':'THR','CSO':'CYS','SEP':'SER','M3L':'LYS',\n",
        "          'HSK':'HIS','SAC':'SER','PCA':'GLU','DAL':'ALA',\n",
        "          'CME':'CYS','CSD':'CYS','OCS':'CYS','DPR':'PRO',\n",
        "          'B3K':'LYS','ALY':'LYS','YCM':'CYS','MLZ':'LYS',\n",
        "          '4BF':'TYR','KCX':'LYS','B3E':'GLU','B3D':'ASP',\n",
        "          'HZP':'PRO','CSX':'CYS','BAL':'ALA','HIC':'HIS',\n",
        "          'DBZ':'ALA','DCY':'CYS','DVA':'VAL','NLE':'LEU',\n",
        "          'SMC':'CYS','AGM':'ARG','B3A':'ALA','DAS':'ASP',\n",
        "          'DLY':'LYS','DSN':'SER','DTH':'THR','GL3':'GLY',\n",
        "          'HY3':'PRO','LLP':'LYS','MGN':'GLN','MHS':'HIS',\n",
        "          'TRQ':'TRP','B3Y':'TYR','PHI':'PHE','PTR':'TYR',\n",
        "          'TYS':'TYR','IAS':'ASP','GPL':'LYS','KYN':'TRP',\n",
        "          'CSD':'CYS','SEC':'CYS'}\n",
        "\n",
        "def pdb_to_string(pdb_file, chains=None, models=[1]):\n",
        "  '''read pdb file and return as string'''\n",
        "\n",
        "  if chains is not None:\n",
        "    if \",\" in chains: chains = chains.split(\",\")\n",
        "    if not isinstance(chains,list): chains = [chains]\n",
        "  if models is not None:\n",
        "    if not isinstance(models,list): models = [models]\n",
        "\n",
        "  modres = {**MODRES}\n",
        "  lines = []\n",
        "  seen = []\n",
        "  model = 1\n",
        "  for line in open(pdb_file,\"rb\"):\n",
        "    line = line.decode(\"utf-8\",\"ignore\").rstrip()\n",
        "    if line[:5] == \"MODEL\":\n",
        "      model = int(line[5:])\n",
        "    if models is None or model in models:\n",
        "      if line[:6] == \"MODRES\":\n",
        "        k = line[12:15]\n",
        "        v = line[24:27]\n",
        "        if k not in modres and v in residue_constants.restype_3to1:\n",
        "          modres[k] = v\n",
        "      if line[:6] == \"HETATM\":\n",
        "        k = line[17:20]\n",
        "        if k in modres:\n",
        "          line = \"ATOM  \"+line[6:17]+modres[k]+line[20:]\n",
        "      if line[:4] == \"ATOM\":\n",
        "        chain = line[21:22]\n",
        "        if chains is None or chain in chains:\n",
        "          atom = line[12:12+4].strip()\n",
        "          resi = line[17:17+3]\n",
        "          resn = line[22:22+5].strip()\n",
        "          if resn[-1].isalpha(): # alternative atom\n",
        "            resn = resn[:-1]\n",
        "            line = line[:26]+\" \"+line[27:]\n",
        "          key = f\"{model}_{chain}_{resn}_{resi}_{atom}\"\n",
        "          if key not in seen: # skip alternative placements\n",
        "            lines.append(line)\n",
        "            seen.append(key)\n",
        "      if line[:5] == \"MODEL\" or line[:3] == \"TER\" or line[:6] == \"ENDMDL\":\n",
        "        lines.append(line)\n",
        "  return \"\\n\".join(lines)\n",
        "\n",
        "def relax_me(pdb_in, pdb_out):\n",
        "  pdb_str = pdb_to_string(pdb_in)\n",
        "  protein_obj = protein.from_pdb_string(pdb_str)\n",
        "  amber_relaxer = relax.AmberRelaxation(\n",
        "    max_iterations=0,\n",
        "    tolerance=2.39,\n",
        "    stiffness=10.0,\n",
        "    exclude_residues=[],\n",
        "    max_outer_iterations=3,\n",
        "    use_gpu=DEVICE == \"GPU\")\n",
        "  relaxed_pdb_lines, _, _ = amber_relaxer.process(prot=protein_obj)\n",
        "  with open(pdb_out, 'w') as f:\n",
        "      f.write(relaxed_pdb_lines)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "I51OdQIa7Xuz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title #RELAX\n",
        "#@markdown - Click the little ▶ play icon to the left to get an upload prompt.\n",
        "#@markdown - After relax is done, `relaxed.pdb` will automatically download.\n",
        "#@markdown - If download was blocked, click the little folder 📁 icon on the left, right-click `relaxed.pdb` and select Download!\n",
        "\n",
        "from google.colab import files\n",
        "pdb_dict = files.upload()\n",
        "relax_me(pdb_in=list(pdb_dict.keys())[0],\n",
        "         pdb_out=\"relaxed.pdb\")\n",
        "files.download(f'relaxed.pdb')"
      ],
      "metadata": {
        "id": "rEekqQXQ-Pec",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}